//      kmean.h
//      
//      K-mean algorithm and basic template for an Algorithm class
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

#ifndef __KNEAN_H
#define __KMEAN_H
#include "stl-helper.h"
/** Basic members of an Algorithm class
*    constructor(<variable_map>, <data>) // to pass parameters
*    save(<string>) // output result
*    run()
*
*/
template<typename D_>
class KMean
{
public:
  typedef D_ data_t;
  typedef typename D_::datum_t datum_t;
  typedef typename D_::datum_t::x_type x_type;
  typedef std::vector<x_type> result_t;
public:
  KMean(boost::program_options::variables_map& vm, data_t& data);
  void run();

  const result_t& result() const { return c_; }
  
  void save(const std::string& fname) const;
protected:
  void init_centers();
  int find_minc(const x_type& x, double& min_dist) const;
  void debug_print(std::ostream& debug, int iter, double total_dist);
  void assign_to_centers(double &total_dist);
  void update_centers();

  boost::program_options::variables_map& vm_;
  data_t& data_;

  /** the number of clusters */
  int K_;
  /** maximum number of iterations */
  int maxiter_;
  /** the centers of the clusters */
  result_t c_;
  
  /** Accumulation variables */
  std::vector<x_type> ACC_center;
  vi_t ACC_count;
};


// ============= IMPLEMENTATION ===============
template <typename D>
KMean<D>::KMean(boost::program_options::variables_map& vm, data_t& data)
  : vm_(vm), data_(data)
{
  K_ = vm_["K"].as<int>();
  maxiter_ = vm_["iter"].as<int>();
}

template <typename D>
void KMean<D>::run()
{
  std::string debugFile = vm_["debug"].as<std::string>();
  FileAccessWrite(debug, debugFile);
  
  init_centers();
  
  ACC_center = c_;
  ACC_count = vi_t(K_, 0);

  double total_dist = 1e-6;
  for (int iter = 0; iter < maxiter_; iter++) {
    // reset accumulation variables
    for (int c = 0; c < K_; c++) {
      ACC_center.at(c).zeros();
      ACC_count.at(c) = 0;
    }

    // assign points to closest centers
    double old_total = total_dist;
    assign_to_centers(total_dist);

    debug_print(debug, iter, total_dist);

    // check termination conditions
    //MSG(total_dist << " " << old_total << " c = " << fabs((old_total - total_dist) / old_total));
    if (fabs((old_total - total_dist) / old_total) < 1e-2) {
      break;
    }

    // update centers
    update_centers();

  }
  debug.close();
}

template <typename D>
void KMean<D>::save(const std::string& fname) const
{
  std::ofstream out;
  FileAccess::open2write(out, fname);
  //out << sz(c_) << "\t" << (sz(c_) > 0 ? sz(c_[0]) : 0) << "\t0" << ENDL;
  BOOST_FOREACH(const x_type& center, c_) {
    datum_t d(center, 0);
    d.save(out, false, "\t");
    out << ENDL;
  }
  out.close();
}

template <typename D>
void KMean<D>::init_centers()
{
  // init centers from the data points
  c_.clear();
  int n = data_.size();
  vi_t idx = RAND::randperm(n);
  for (int c = 0; c < K_; c++) {
    int r = rand() % (n-c)+c;   // a random index from remaining indices
    std::swap(idx[r], idx[c]);  // swap to current center index
    c_.pb(data_(idx[c]).x());
  }
}

template <typename D>
int KMean<D>::find_minc(const x_type& x_i, double& min_dist) const
{
  int min_c = -1;
  for (int c = 0; c < K_; c++) {
    double dist = arma::norm(c_.at(c) - x_i, 2);
    if (dist < min_dist) {
      min_c = c;
      min_dist = dist;
    }
  }
  return min_c;
}

template <typename D>
void KMean<D>::debug_print(std::ostream& debug, int iter, double total_dist)
{
  MSG("iter = " << iter << " total_dist = " << total_dist);
  debug << iter << " " << total_dist << ENDL;
  for (int c = 0; c < K_; c++) {
    debug << c_[c].t();
  }
}

template <typename D>
void KMean<D>::update_centers()
{
  for (int c = 0; c < K_; c++) {
    if (ACC_count.at(c) != 0)
      c_.at(c) = ACC_center.at(c) / ACC_count.at(c);
    else { // no point assign to cluster c, assign a random point
      int r = rand() % data_.size();
      c_.at(c) = data_(r).x();
    }
  }
}

template <typename D>
void KMean<D>::assign_to_centers(double &total_dist)
{
  total_dist = 0;
  int n = data_.size();
  for (int i = 0; i < n; i++) {
    const x_type& x_i = data_(i).x();
    double min_dist = DOUBLE_INFINITY;
    int min_c = find_minc(x_i, min_dist);

    total_dist += min_dist;
    ACC_center.at(min_c) += x_i;
    ACC_count.at(min_c) ++;
  }
}

#endif
